{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Token Counting in ZhipuAI GLM API\n",
    "\n",
    "**This tutorial is available in English and is attached below the Chinese explanation**\n",
    "\n",
    "此代码将讲述在 GLM 模型 是如何计算 Token 的消耗的。这将能帮你对在 API 使用中的 Token 计算产生更清晰的认识。\n",
    "\n",
    "This cookbook will describe how Token consumption is calculated in the GLM model. This will help you have a clearer understanding of Token calculations in API usage."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "71c55dc5a09277b"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load Tokenizer and count\n",
    "\n",
    "由于没有任何的公开的 Zhipu AI token计算工具，因此，我使用了 [chatglm3-6b](https://huggingface.co/THUDM/chatglm3-6b) 这个开源模型的 tokenizer进行加载。这种计算方式仅能作为参考，尚且不能认定是最终的 API token 计算方式。具体的计价方式以官方文档为主。\n",
    "\n",
    "Since there is no public Zhipu AI token calculation tool, I used the tokenizer of this open source model [chatglm3-6b](https://huggingface.co/THUDM/chatglm3-6b) to load it. This calculation method can only be used as a reference and cannot be considered as the final API token calculation method."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8977deb25864ec9f"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"chatglm3-6b\", trust_remote_code=True, encode_special_tokens=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:51:13.105834Z",
     "start_time": "2024-01-24T12:51:11.590578Z"
    }
   },
   "id": "e463ead175234bf7"
  },
  {
   "cell_type": "markdown",
   "source": [
    "我们书写一个Python 代码，分别有以下两个功能\n",
    "\n",
    "1. `count_encode` : 将自然语言转换为 token 并计算 token 消耗。\n",
    "2. `decode`: 输入一串 token 进行解码，并展示 词语是怎么被拆开的。\n",
    "\n",
    "we write a Python code with the following two functions:\n",
    "\n",
    "1. `count_encode`: Convert natural language into token and calculate token consumption.\n",
    "2. `decode`: Input a string of tokens to decode and show how the words are split."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7125d5fca775b004"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count Token: 7\n",
      "Token IDs: [64790, 64792, 985, 323, 260, 2254, 16948]\n",
      "Decode sentence: ['[gMASK]', 'sop', '▁here', '▁is', '▁a', '▁text', '▁demo']\t"
     ]
    }
   ],
   "source": [
    "def count_encode(inputs: str = \"\"):\n",
    "    encoded_input = tokenizer.encode(inputs)\n",
    "    num_tokens = len(encoded_input)\n",
    "    return encoded_input, num_tokens\n",
    "\n",
    "\n",
    "def decode(inputs: list = []):\n",
    "    decode_sentence = tokenizer.convert_ids_to_tokens(inputs)\n",
    "    return decode_sentence\n",
    "\n",
    "\n",
    "text = \"here is a text demo\"\n",
    "encoded_input, num_tokens = count_encode(text)\n",
    "print(f\"Count Token: {num_tokens}\")\n",
    "print(f\"Token IDs: {encoded_input}\")\n",
    "decode_sentence = decode(encoded_input)\n",
    "print(f\"Decode sentence: {decode_sentence}\", end=\"\\t\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:51:13.111618Z",
     "start_time": "2024-01-24T12:51:13.107437Z"
    }
   },
   "id": "5b7ed321306f5849"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Token calculation in conversation\n",
    "\n",
    "刚才的代码并不是在真正对话中 token 计算的使用量， 在对话中，由于存在模型的 **对话模板** 和 **聊天历史**，对话中的 token 计算会明显变高。\n",
    "\n",
    "我先展示了在对话场景中，一条 message 的实际对应的token数量。\n",
    "\n",
    "The code just now is not the usage of token calculation in the real conversation. In the conversation, due to the existence of the model's **conversation template** and **chat history**, the token calculation in the conversation will be significantly higher.\n",
    "\n",
    "I first showed the actual number of tokens corresponding to a message in the conversation scenario."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2f78f812a8c1f0ea"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count Token: 11\n",
      "Token IDs: [64790, 64792, 64795, 265, 13, 985, 323, 260, 2254, 16948, 64796]\n",
      "['[gMASK]', 'sop', '<|user|>', '▁▁', '<0x0A>', '▁here', '▁is', '▁a', '▁text', '▁demo', '<|assistant|>']\t"
     ]
    }
   ],
   "source": [
    "messages = []\n",
    "messages.append({\"role\": \"user\", \"content\": \"here is a text demo\"})\n",
    "chat_inputs = tokenizer.apply_chat_template(messages,\n",
    "                                            add_generation_prompt=True,\n",
    "                                            tokenize=True)\n",
    "num_tokens = len(chat_inputs)\n",
    "print(f\"Count Token: {num_tokens}\")\n",
    "print(f\"Token IDs: {chat_inputs}\")\n",
    "model_input_tokens = tokenizer.convert_ids_to_tokens(chat_inputs)\n",
    "print(model_input_tokens, end=\"\\t\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:51:13.140806Z",
     "start_time": "2024-01-24T12:51:13.110979Z"
    }
   },
   "id": "af6673d5c30f8841"
  },
  {
   "cell_type": "markdown",
   "source": [
    "可以看到，由于 **对话模板** 的存在，这次的 token 计算 多出了 4 个 token 数量，这个数量就是模板的占用。**这部分 token 会在 API 进行计费的过程中被扣除**。\n",
    "\n",
    "It can be seen that due to the existence of **dialogue template**, the token calculation this time has 4 more tokens, which is the occupation of the template. **This part of the token will be deducted during the API billing process**."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7d59444531c388c6"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Calculation of token consumption in actual conversations\n",
    "\n",
    "现在，我们开始模拟在实际的对话中，token 消耗的方式，首先，你需要填写 API \n",
    "Now, we start to simulate the way token is consumed in an actual conversation. First, you need to fill in the API"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f2b9d929eee22dfd"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "import os\n",
    "from zhipuai import ZhipuAI\n",
    "\n",
    "os.environ[\"ZHIPUAI_API_KEY\"] = \"your api key\""
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:51:13.261739Z",
     "start_time": "2024-01-24T12:51:13.131065Z"
    }
   },
   "id": "1293b5a0561a0cfb"
  },
  {
   "cell_type": "markdown",
   "source": [
    "接着，我们使用一段带有历史记录的对话来计算随着 历史记录的不断增加， Token的消耗速度。\n",
    "我书写了一段Python脚本，在这里，你可以和模型进行连续对话，代码将会返回随着对话的进行，每次输入的 Token 数量 和 输出的 Token数量。\n",
    "这个代码模拟了一个简单的命令行交互画面，在这里你可以和大模型进行交互，同时，输出会计算 你的 输入 token 消耗 和 模型的输出token消耗。\n",
    "请注意，`输入消耗` 在 API 消耗的值，是 用户当前输入和对话历史的全部内容，因此，你将会在对话的过程看到 输入使用的 token 逐渐增加。\n",
    "\n",
    "Next, we use a conversation with historical records to calculate the Token consumption rate as the historical records continue to increase.\n",
    "I wrote a Python script, where you can have a continuous conversation with the model, and the code will return the number of Tokens input and the number of Tokens output each time as the conversation proceeds.\n",
    "\n",
    "This code simulates a simple command line interaction screen, where you can interact with the large model. At the same time, the output will calculate your input token consumption and the model's output token consumption.\n",
    "\n",
    "Please note that the value consumed by `input consumption` in the API is the entire content of the user's current input and conversation history. Therefore, you will see the token used for input gradually increase during the conversation."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "187bed9b6a2e5caa"
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "class ChatSession:\n",
    "    def __init__(self, max_tokens=8192):\n",
    "        self.client = ZhipuAI()\n",
    "        self.history = []\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(\n",
    "            \"chatglm3-6b\",\n",
    "            trust_remote_code=True,\n",
    "            encode_special_tokens=True\n",
    "        )\n",
    "        self.max_tokens = max_tokens\n",
    "\n",
    "    def calculate_tokens(self, inputs):\n",
    "        return len(tokenizer.apply_chat_template(\n",
    "            inputs,\n",
    "            add_generation_prompt=True,\n",
    "            tokenize=True)\n",
    "        )\n",
    "\n",
    "    def chat(self):\n",
    "        while True:\n",
    "            user_input = input(\"You: \")\n",
    "            if user_input.lower() in [\"exit\", \"quit\"]:\n",
    "                break\n",
    "            self.history.append({\"role\": \"user\", \"content\": user_input})\n",
    "            print(f\"Human: {user_input}\")\n",
    "            user_tokens = self.calculate_tokens(self.history)\n",
    "            print(\"======================\")\n",
    "            print(f\"Human message with history tokens: {user_tokens}\")\n",
    "\n",
    "            response = self.client.chat.completions.create(\n",
    "                model=\"glm-4\",\n",
    "                messages=self.history,\n",
    "                top_p=0.7,\n",
    "                temperature=0.9,\n",
    "                stream=False,\n",
    "                max_tokens=self.max_tokens,\n",
    "            )\n",
    "\n",
    "            reply = response.choices[0].message.content\n",
    "            print(f\"AI: {reply}\")\n",
    "            self.history.append({\"role\": \"assistant\", \"content\": reply})\n",
    "            reply_tokens = self.calculate_tokens(self.history[-2:])\n",
    "            print(\"======================\")\n",
    "            print(f\"AI response tokens: {reply_tokens}\", end=\"\\n**********************\\n\")\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:51:13.267973Z",
     "start_time": "2024-01-24T12:51:13.265482Z"
    }
   },
   "id": "f03d04d64937fc37"
  },
  {
   "cell_type": "markdown",
   "source": [
    "开始对话吧，通过这次对话，你应该会意识到缩短历史记录，并定期清空历史记录的重要性。\n",
    "\n",
    "Start a conversation！That should help you realize the importance of shortening your history and clearing it regularly."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "92bda16255169230"
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: hello! what a nice day!\n",
      "======================\n",
      "Human message with history tokens: 13\n",
      "AI: Hello! Indeed, it looks like a beautiful day! Is there anything I can help you with today?\n",
      "======================\n",
      "AI response tokens: 37\n",
      "**********************\n",
      "Human: Are you good at joke?\n",
      "======================\n",
      "Human message with history tokens: 46\n",
      "AI: I can certainly share a joke or two! Here's a light-hearted one for you:\n",
      "\n",
      "Why don't scientists trust atoms?\n",
      "\n",
      "Because they make up everything!\n",
      "\n",
      "Feel free to share a joke with me too, or let me know if you'd like to hear another one.\n",
      "======================\n",
      "AI response tokens: 80\n",
      "**********************\n",
      "Human: give me a example\n",
      "======================\n",
      "Human message with history tokens: 121\n",
      "AI: Sure! Here's another one:\n",
      "\n",
      "Why did the scarecrow win an award?\n",
      "\n",
      "Because he was outstanding in his field!\n",
      "\n",
      "And here's a classic math joke:\n",
      "\n",
      "What did the zero say to the eight?\n",
      "\n",
      "\"Nice belt!\" 🙂\n",
      "\n",
      "Let me know if you want more jokes or if there's anything else I can assist you with!\n",
      "======================\n",
      "AI response tokens: 95\n",
      "**********************\n",
      "Human: tell me some joke about school or university?\n",
      "======================\n",
      "Human message with history tokens: 218\n",
      "AI: Certainly! Here are a couple of jokes related to school and university:\n",
      "\n",
      "1. University Joke:\n",
      "Why don't scientists trust atoms?\n",
      "\n",
      "Because they make up everything, including the grades!\n",
      "\n",
      "2. School Joke:\n",
      "Why was the math book sad?\n",
      "\n",
      "Because it had too many problems.\n",
      "\n",
      "And here's another one:\n",
      "\n",
      "Why did the student bring a sunflower to the computer class?\n",
      "\n",
      "To brighten up the byte!\n",
      "\n",
      "Hope these jokes bring a smile to your face. If you want more, just let me know!\n",
      "======================\n",
      "AI response tokens: 142\n",
      "**********************\n",
      "Human: Thank you!\n",
      "======================\n",
      "Human message with history tokens: 351\n",
      "AI: You're welcome! I'm glad I could bring a smile to your day. If you have any more questions or need assistance with anything else, feel free to ask. Have a great day!\n",
      "======================\n",
      "AI response tokens: 53\n",
      "**********************\n"
     ]
    }
   ],
   "source": [
    "chat_session = ChatSession()\n",
    "chat_session.chat()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-01-24T12:52:19.532008Z",
     "start_time": "2024-01-24T12:51:13.267799Z"
    }
   },
   "id": "a651bbaad4bd2512"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
